package com.ksyun.campus.dataserver.services;

import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.stereotype.Component;
import com.fasterxml.jackson.databind.ObjectMapper;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import java.io.File;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
@Component
public class RegistService implements ApplicationRunner {

    @Value("${dataserver.node-id:data-server-1}")
    private String nodeId;
    
    @Value("${dataserver.zookeeper.connect-string:localhost:2181}")
    private String zkConnectString;
    
    @Value("${dataserver.zookeeper.session-timeout:30000}")
    private int sessionTimeout;
    
    @Value("${dataserver.zookeeper.connection-timeout:15000}")
    private int connectionTimeout;
    
    @Value("${dataserver.heartbeat-interval:10000}")
    private int heartbeatInterval;
    
    @Value("${dataserver.heartbeat.retry-interval:5000}")
    private int heartbeatRetryInterval;
    
    @Value("${dataserver.heartbeat.max-retries:3}")
    private int maxHeartbeatRetries;
    
    @Value("${server.port:9001}")
    private int serverPort;
    
    @Value("${dataserver.data.path:/data}")
    private String dataPath;
    
    @Autowired
    private CuratorFramework zkClient;
    private ScheduledExecutorService heartbeatExecutor;
    private String nodePath;
    private String heartbeatPath;
    private final ObjectMapper objectMapper;
    private final AtomicBoolean isRunning;
    private final AtomicBoolean isRegistered;
    
    public RegistService() {
        this.objectMapper = new ObjectMapper();
        this.isRunning = new AtomicBoolean(false);
        this.isRegistered = new AtomicBoolean(false);
    }
    
    @PostConstruct
    public void init() {
        // 统一使用配置类提供的 CuratorFramework（无 namespace），路径使用绝对路径 /minfs/...
        nodePath = "/minfs/dataservers/" + nodeId;
        heartbeatPath = "/minfs/dataservers/heartbeat/" + nodeId;
        log.info("DataServer initialized with nodeId: {}, port: {}, dataPath: {}", 
                nodeId, serverPort, dataPath);
    }
    
    // CuratorFramework 由配置类创建与启动，无需在此重复初始化
    
    @Override
    public void run(ApplicationArguments args) throws Exception {
        try {
            registToCenter();
            startHeartbeat();
            isRunning.set(true);
            log.info("DataServer registration and heartbeat service started successfully");
        } catch (Exception e) {
            log.error("Failed to start DataServer services", e);
            throw e;
        }
    }
    
    /**
     * 注册到ZooKeeper中心
     */
    public void registToCenter() {
        if (isRegistered.get()) {
            log.debug("Already registered to ZooKeeper, skipping registration");
            return;
        }
        
        try {
            // 创建数据服务器节点信息
            Map<String, Object> nodeInfo = createNodeInfo();
            
            // 注册到ZooKeeper
            String nodeData = objectMapper.writeValueAsString(nodeInfo);
            
            // 创建持久节点
            if (zkClient.checkExists().forPath(nodePath) == null) {
                zkClient.create().creatingParentsIfNeeded().forPath(nodePath, nodeData.getBytes(StandardCharsets.UTF_8));
                log.info("Successfully registered to ZooKeeper: {}", nodePath);
            } else {
                zkClient.setData().forPath(nodePath, nodeData.getBytes(StandardCharsets.UTF_8));
                log.info("Successfully updated registration in ZooKeeper: {}", nodePath);
            }
            
            // 创建心跳节点（临时节点）
            createHeartbeatNode();
            
            isRegistered.set(true);
            
        } catch (Exception e) {
            log.error("Failed to register to ZooKeeper", e);
            isRegistered.set(false);
            throw new RuntimeException("Registration failed", e);
        }
    }
    
    /**
     * 创建节点信息
     */
    private Map<String, Object> createNodeInfo() {
        Map<String, Object> nodeInfo = new HashMap<>();
        // 使用dsId字段名以匹配DataServerInfo类
        nodeInfo.put("dsId", nodeId);
        nodeInfo.put("host", "localhost");
        nodeInfo.put("port", serverPort);
        nodeInfo.put("status", "ACTIVE");
        java.time.format.DateTimeFormatter formatter = java.time.format.DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
        String currentTime = java.time.LocalDateTime.now().format(formatter);
        nodeInfo.put("lastSeen", currentTime);
        nodeInfo.put("lastHeartbeat", currentTime);
        
        // 计算磁盘容量和使用情况
        Map<String, Object> capacityInfo = getCapacityInfo();
        nodeInfo.putAll(capacityInfo);
        
        return nodeInfo;
    }
    
    /**
     * 创建心跳节点
     */
    private void createHeartbeatNode() {
        try {
            if (zkClient.checkExists().forPath(heartbeatPath) == null) {
                zkClient.create()
                    .creatingParentsIfNeeded()
                    .withMode(org.apache.zookeeper.CreateMode.EPHEMERAL)
                    .forPath(heartbeatPath, "heartbeat".getBytes(StandardCharsets.UTF_8));
                log.info("Successfully created heartbeat node: {}", heartbeatPath);
            }
        } catch (Exception e) {
            log.error("Failed to create heartbeat node", e);
            throw new RuntimeException("Heartbeat node creation failed", e);
        }
    }
    
    /**
     * 获取容量信息，改进错误处理
     */
    private Map<String, Object> getCapacityInfo() {
        Map<String, Object> capacityInfo = new HashMap<>();
        try {
            Path dataDir = Paths.get(dataPath);
            if (Files.exists(dataDir)) {
                File dataFile = dataDir.toFile();
                long totalSpace = dataFile.getTotalSpace();
                long freeSpace = dataFile.getFreeSpace();
                long usedSpace = totalSpace - freeSpace;
                // 使用DataServerInfo类期望的字段名
                capacityInfo.put("capacityBytes", totalSpace);
                capacityInfo.put("usedBytes", usedSpace);
                // 保留原有字段以兼容其他用途
                capacityInfo.put("totalCapacity", totalSpace);
                capacityInfo.put("freeCapacity", freeSpace);
                capacityInfo.put("usedCapacity", usedSpace);
                capacityInfo.put("usagePercent", (double) usedSpace / totalSpace * 100);
                
                // 计算数据文件大小
                long dataSize = calculateDataSize(dataPath);
                capacityInfo.put("dataSize", dataSize);
                
                log.debug("Capacity info - Total: {}MB, Free: {}MB, Used: {}MB, Data: {}MB", 
                        totalSpace / (1024 * 1024), 
                        freeSpace / (1024 * 1024), 
                        usedSpace / (1024 * 1024),
                        dataSize / (1024 * 1024));
            } else {
                log.warn("Data directory does not exist: {}", dataPath);
                capacityInfo.put("capacityBytes", 0L);
                capacityInfo.put("usedBytes", 0L);
                capacityInfo.put("totalCapacity", 0L);
                capacityInfo.put("freeCapacity", 0L);
                capacityInfo.put("usedCapacity", 0L);
                capacityInfo.put("usagePercent", 0.0);
                capacityInfo.put("dataSize", 0L);
            }
        } catch (Exception e) {
            log.error("Error getting capacity info", e);
            // 设置默认值
            capacityInfo.put("capacityBytes", -1L);
            capacityInfo.put("usedBytes", -1L);
            capacityInfo.put("totalCapacity", -1L);
            capacityInfo.put("freeCapacity", -1L);
            capacityInfo.put("usedCapacity", -1L);
            capacityInfo.put("usagePercent", -1.0);
            capacityInfo.put("dataSize", -1L);
        }
        return capacityInfo;
    }
    
    /**
     * 计算数据目录大小，改进错误处理
     */
    private long calculateDataSize(String path) {
        try {
            Path dataDir = Paths.get(path);
            if (!Files.exists(dataDir)) {
                return 0;
            }
            
            return Files.walk(dataDir)
                    .filter(Files::isRegularFile)
                    .mapToLong(p -> {
                        try {
                            return Files.size(p);
                        } catch (Exception e) {
                            log.debug("Failed to get size for file: {}", p, e);
                            return 0;
                        }
                    })
                    .sum();
        } catch (Exception e) {
            log.error("Error calculating data size for path: {}", path, e);
            return 0;
        }
    }
    
    /**
     * 启动心跳服务，改进错误处理和重试机制
     */
    private void startHeartbeat() {
        heartbeatExecutor = Executors.newSingleThreadScheduledExecutor(r -> {
            Thread t = new Thread(r, "heartbeat-thread-" + nodeId);
            t.setDaemon(true);
            return t;
        });
        
        heartbeatExecutor.scheduleAtFixedRate(() -> {
            if (!isRunning.get()) {
                return;
            }
            
            try {
                sendHeartbeat();
            } catch (Exception e) {
                log.error("Error sending heartbeat", e);
                handleHeartbeatError();
            }
        }, 0, heartbeatInterval, TimeUnit.MILLISECONDS);
        
        log.info("Started heartbeat service with interval: {}ms", heartbeatInterval);
    }
    
    /**
     * 发送心跳，改进重试机制
     */
    private void sendHeartbeat() {
        int retryCount = 0;
        while (retryCount < maxHeartbeatRetries) {
            try {
                // 更新心跳时间戳
                Map<String, Object> heartbeatInfo = createHeartbeatInfo();
                String heartbeatData = objectMapper.writeValueAsString(heartbeatInfo);
                
                // 更新心跳节点
                if (zkClient.checkExists().forPath(heartbeatPath) != null) {
                    zkClient.setData().forPath(heartbeatPath, heartbeatData.getBytes(StandardCharsets.UTF_8));
                } else {
                    // 如果心跳节点不存在，重新创建
                    createHeartbeatNode();
                }
                
                log.debug("Heartbeat sent successfully: {}", heartbeatPath);
                return;
                
            } catch (Exception e) {
                retryCount++;
                log.warn("Heartbeat attempt {} failed: {}", retryCount, e.getMessage());
                
                if (retryCount < maxHeartbeatRetries) {
                    try {
                        Thread.sleep(heartbeatRetryInterval);
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                        return;
                    }
                }
            }
        }
        
        // 所有重试都失败了
        log.error("All heartbeat attempts failed after {} retries", maxHeartbeatRetries);
        throw new RuntimeException("Heartbeat failed after maximum retries");
    }
    
    /**
     * 创建心跳信息
     */
    private Map<String, Object> createHeartbeatInfo() {
        Map<String, Object> heartbeatInfo = new HashMap<>();
        heartbeatInfo.put("timestamp", System.currentTimeMillis());
        heartbeatInfo.put("status", "alive");
        heartbeatInfo.put("nodeId", nodeId);
        
        // 更新容量信息
        Map<String, Object> capacityInfo = getCapacityInfo();
        heartbeatInfo.putAll(capacityInfo);
        
        return heartbeatInfo;
    }
    
    /**
     * 处理心跳错误
     */
    private void handleHeartbeatError() {
        try {
            // 尝试重新注册
            if (!isRegistered.get()) {
                registToCenter();
            }
            
            // 尝试重新创建心跳节点
            createHeartbeatNode();
            
        } catch (Exception re) {
            log.error("Failed to recover from heartbeat error", re);
            // 标记为未注册状态，下次心跳会尝试重新注册
            isRegistered.set(false);
        }
    }
    
    /**
     * 获取数据服务器列表
     */
    public List<Map<String, Integer>> getDslist() {
        // TODO: 从ZooKeeper获取所有数据服务器列表
        // 这里可以实现从ZK获取所有DataServer的信息
        return null;
    }
    
    /**
     * 检查服务状态
     */
    public boolean isHealthy() {
        return isRunning.get() && isRegistered.get() && 
               zkClient != null && zkClient.getZookeeperClient().isConnected();
    }
    
    /**
     * 获取服务状态信息
     */
    public Map<String, Object> getStatus() {
        Map<String, Object> status = new HashMap<>();
        status.put("running", isRunning.get());
        status.put("registered", isRegistered.get());
        status.put("zkConnected", zkClient != null && zkClient.getZookeeperClient().isConnected());
        status.put("nodeId", nodeId);
        status.put("port", serverPort);
        status.put("uptime", System.currentTimeMillis() - getStartTime());
        
        return status;
    }
    
    private long startTime = System.currentTimeMillis();
    
    private long getStartTime() {
        return startTime;
    }
    
    @PreDestroy
    public void cleanup() {
        log.info("Starting DataServer cleanup...");
        
        isRunning.set(false);
        
        try {
            if (heartbeatExecutor != null) {
                heartbeatExecutor.shutdown();
                if (!heartbeatExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                    heartbeatExecutor.shutdownNow();
                }
            }
            
            if (zkClient != null) {
                zkClient.close();
            }
            
            log.info("DataServer cleanup completed successfully");
            
        } catch (Exception e) {
            log.error("Error during cleanup", e);
        }
    }
}


